﻿//************************************************************************************************
// Copyright © 2003-2012 Steven M. Cohn. All Rights Reserved.
//
// Facility:	Microsoft Development Environment 2012  Version 11.0.50727.1
// Environment:	Microsoft .NET Framework 4.5            Version v4.0.30319
//
// Description: Provides an enumerable stream of index.dat cached cookies.
//
// Revision History:
// -Who--------------- -When---------- -What------------------------------------------------------ 
// Steven M. Cohn		18-Oct-2003		New
// Steven M. Cohn		12-Oct-2004		VS.NET 2005, .NET 2.0
// Steven M. Cohn		20-Nov-2012		Refresh! VS2010 .NET 4.5 64-bit
//************************************************************************************************

namespace CookieChipper
{
	using System;
	using System.Collections.Generic;
	using System.Collections.Specialized;
	using System.IO;
	using System.Security.Principal;
	using System.Text;


	/// <summary>
	/// Provides an enumerable stream of index.dat cached cookies.
	/// </summary>

	internal class CookieCacheProvider : IDisposable
	{
		private const string CacheFilename = "index.dat";
		private const string DateFormat = "MM/dd/yyyy h:mm tt";

		private const int BlockSize = 0x80;				//   128 record may contain >=1 blocks

		private const string HeaderSignature = "Client UrlCache MMF Ver 5.2";
		private const int HeaderLen = 72;				//    72 file header record length
		private const int HeaderSignatureLen = 28;
		private const int HeaderFizeSizeOffset = 28;
		private const int HeaderHashOffset = 32;
		private const int HeaderBlocksOffset = 36;
		private const int HeaderAllocBlocksOffset = 40;
		private const int HeaderQuotaOffset = 48;
		private const int HeaderCacheSizeOffset = 56;

		private const string HashSignature = "HASH";
		private const int HashOffset = 0x4000;			// 16384 offset in file to first hash table
		private const int HashHeaderLen = 16;
		private const int HashRecordLen = 8;			// two integers
		private const int HashSignatureLen = 4;
		private const uint HashValid = 0x0000000F;		// inverted low byte (compare with "&")
		private const uint HashREDR = 0x00000005;		// mask

		private const string RecordSignature = "URL ";	// preamble of cache record
		private const int RecordBlocks = 4;				// offset to block count
		private const int RecordFilename = 60;			// offset to filename string
		private const int RecordAccessed = 8;			// offset to created datetime
		private const int RecordUpdated = 16;			// offset to updated datetime
		private const int RecordFileSizeLo = 32;		// offset to low int of file size
		private const int RecordFileSizeHi = 36;		// offset to high int of file size
		private const int RecordDomain = 52;			// offset to domain string
		private const int RecordHits = 84;				// offset to hits count

		// MSIE7 zero pattern read backwards is "0badf00d"
		// MSIE8 zero pattern read backwards is "deadbeef"
		private readonly byte[] Zero7 = new byte[4] { 0x0D, 0xF0, 0xAD, 0x0B };
		private readonly byte[] Zero8 = new byte[4] { 0xEF, 0xBE, 0xAD, 0xDE };

		private string path;
		private string root;
		private FileStream stream;

		private CacheHeader header;
		private List<HashRef> hashtable;
		private StringCollection rogues;
		private int hashindex;

		private bool isDisposed;


		//========================================================================================
		// Lifecycle
		//========================================================================================

		/// <summary>
		/// 
		/// </summary>

		public CookieCacheProvider ()
			: this(Path.Combine(GetDirectoryPath(), CacheFilename))
		{
		}


		/// <summary>
		/// For unit testing, redirect to alternate index.dat path
		/// </summary>
		/// <param name="path"></param>

		public CookieCacheProvider (string path)
		{
			this.path = path;
			header = new CacheHeader(this.path);

			if (!File.Exists(this.path))
			{
				throw new FileNotFoundException(this.path);
			}

			// directory path without index.dat filename
			root = Path.GetDirectoryName(this.path);

			hashindex = 0;
			isDisposed = false;
		}


		/// <summary>
		/// Close and dispose the reader.
		/// </summary>

		public void Dispose ()
		{
			if (!isDisposed)
			{
				if (stream != null)
				{
					try
					{
						stream.Dispose();
					}
					catch
					{
					}
					finally
					{
						stream = null;
					}
				}

				isDisposed = true;
			}
		}


		/// <summary>
		/// 
		/// </summary>

		public CacheHeader Header
		{
			get { return header; }
		}


		//========================================================================================
		// Clean()
		//========================================================================================

		/// <summary>
		/// Clean the cookie cache, preserving only tagged favorites.
		/// </summary>
		/// <returns>
		/// The total count of cookies removed from the cache.
		/// </returns>

		public int Clean (Func<Cookie, bool> locked)
		{
			int count = 0;

			if (Open())
			{
				long startPosition = stream.Position;

				Cookie cookie;
				while ((cookie = GetNext()) != null)
				{
					if (!locked(cookie))
					{
						if (Delete(cookie))
						{
							CookieFileProvider.Delete(cookie.Path);
							count++;
						}
					}
				}
			}

			return count;
		}


		//========================================================================================
		// Delete()
		//========================================================================================

		/// <summary>
		/// Deletes a single cookie from the cache.
		/// </summary>
		/// <param name="cookie">The cookie instance to delete.</param>
		/// <returns>
		/// True if the cookie cache entry was removed successfully.
		/// </returns>

		public bool Delete (Cookie cookie)
		{
			bool success = false;

			HashRef hash = cookie.Hash;

			// if rogue
			if (hash.HashAddress == 0)
			{
				// it's not in the index so we can't say we failed, right!?
				return true;
			}

			// confirm that the hash record is appropriate, has not changed
			stream.Position = hash.HashAddress;

			byte[] buffer = new byte[HashRecordLen];
			if ((stream.Read(buffer, 0, HashRecordLen) == HashRecordLen))
			{
				int address = BitConverter.ToInt32(buffer, 4);
				if (address == hash.DataAddress)
				{
					// confirm that the cookie record is appropriate, has not changed

					buffer = null;
					if (ReadRecord(hash.DataAddress, out buffer))
					{
						var builder = new StringBuilder();
						int c = BitConverter.ToInt32(buffer, RecordFilename);
						while (buffer[c] != 0)
						{
							builder.Append((char)buffer[c++]);
						}

						if (builder.Length > 0)
						{
							string name = Path.GetFileNameWithoutExtension(builder.ToString());
							if (name.Equals(cookie.Name))
							{
								// return to reference address of hash record and wipe it out!

								stream.Position = hash.HashAddress + 4;
								stream.Write(Zero8, 0, 4);

								success = true;
							}
						}
					}
				}
			}

			return success;
		}


		//========================================================================================
		// GetDirectoryPath()
		//========================================================================================

		/// <summary>
		/// Returns the path of the cookie cache directory.
		/// </summary>
		/// <returns></returns>

		public static string GetDirectoryPath ()
		{
			CacheLocation location;
			return GetDirectoryPath(out location);
		}


		/// <summary>
		/// Returns the path of the cookie cache directory and its location.
		/// </summary>
		/// <param name="location"></param>
		/// <returns></returns>

		public static string GetDirectoryPath (out CacheLocation location)
		{
			// determine if UAC is enabled or not!
			bool isElevated = false;
			using (var identity = WindowsIdentity.GetCurrent())
			{
				if (new WindowsPrincipal(identity).IsInRole(WindowsBuiltInRole.Administrator))
				{
					isElevated = true;
				}
			}

			string path;
			if (isElevated)
			{
				// C:\Users\<username>\AppData\Roaming\Microsoft\Windows\Cookies\index.dat
				path = Path.Combine(Environment.GetFolderPath(
					Environment.SpecialFolder.ApplicationData),
					@"Microsoft\Windows\Cookies");

				location = CacheLocation.Roaming;

				if (!Directory.Exists(path) ||
					!File.Exists(Path.Combine(path, "index.dat")))
				{
					// C:\Users\<username>\AppData\Local\Temp\Cookies\index.dat
					path = Path.Combine(Environment.GetFolderPath(
						Environment.SpecialFolder.LocalApplicationData),
						@"Temp\Cookies");

					location = CacheLocation.Local;
				}
			}
			else
			{
				// C:\Users\<username>\AppData\Roaming\Microsoft\Windows\Cookies\Low\index.dat
				path = Path.Combine(Environment.GetFolderPath(
					Environment.SpecialFolder.ApplicationData),
					@"Microsoft\Windows\Cookies\Low");

				location = CacheLocation.RoamingLow;

				if (!Directory.Exists(path) ||
					!File.Exists(Path.Combine(path, "index.dat")))
				{
					// C:\Users\<username>\AppData\Local\Temp\Cookies\Low\index.dat
					path = Path.Combine(Environment.GetFolderPath(
						Environment.SpecialFolder.LocalApplicationData),
						@"Temp\Cookies\Low");

					location = CacheLocation.LocalLow;
				}
			}

			return path;
		}


		//========================================================================================
		// GetNext()
		//========================================================================================

		/// <summary>
		/// Read and return the next available cookie.
		/// </summary>
		/// <returns>A Cookie.</returns>

		public Cookie GetNext ()
		{
			Cookie cookie = null;

			if ((stream == null) || (hashindex >= hashtable.Count))
			{
				if (rogues.Count > 0)
				{
					cookie = new Cookie(new HashRef(), rogues[0]);
					cookie.Path = Path.Combine(root, cookie.Name + ".txt");
					cookie.FileSize = (int)(new FileInfo(cookie.Path).Length);
					cookie.State = CookieState.Rogue;

					rogues.RemoveAt(0);
				}

				return cookie;
			}

			byte[] buffer;
			HashRef hash = hashtable[hashindex++];

			if (ReadRecord(hash.DataAddress, out buffer))
			{
				// Now that we have the blocks in-memory, we can first skip to the
				// path so we can get the cookie name

				var builder = new StringBuilder();
				int c = BitConverter.ToInt32(buffer, RecordFilename);
				while (buffer[c] != 0)
				{
					builder.Append((char)buffer[c++]);
				}

				if (builder.Length == 0)
				{
					// need to return something to the caller to indicate an empty record using an
					// empty name; this should not happen since we're using the hash table to find
					// valid records!
					return new Cookie(hash, String.Empty);
				}

				string filename = builder.ToString();

				cookie = new Cookie(hash, Path.GetFileNameWithoutExtension(filename));
				cookie.Path = Path.Combine(root, filename);

				// And then we can back up and read other information...

				// timestamps
				cookie.Accessed = DateTime.FromFileTime(BitConverter.ToInt64(buffer, RecordAccessed));
				cookie.Updated = DateTime.FromFileTime(BitConverter.ToInt64(buffer, RecordUpdated));

				// file size
				int lo = BitConverter.ToInt32(buffer, RecordFileSizeLo);
				int hi = BitConverter.ToInt32(buffer, RecordFileSizeHi);
				cookie.FileSize = lo + (hi << 32);

				// domain
				builder = new StringBuilder();
				c = BitConverter.ToInt32(buffer, RecordDomain);
				while (buffer[c] != 0)
				{
					builder.Append((char)buffer[c++]);
				}

				string domain = builder.ToString();
				int amper = domain.IndexOf('@') + 1;
				int length = domain.EndsWith("/") ? domain.Length - 1 : domain.Length;
				cookie.Domain = domain.Substring(amper, length - amper);

				cookie.Hits = BitConverter.ToInt32(buffer, RecordHits);

				if (rogues.Contains(cookie.Name))
				{
					rogues.Remove(cookie.Name);
				}

				return cookie;
			}

			return cookie;
		}


		private bool ReadRecord (long address, out byte[] buffer)
		{
			if (stream.Seek(address, SeekOrigin.Begin) != address)
			{
				throw new FileLoadException("Could not find start of record");
			}

			byte[] head = new byte[BlockSize];

			// read 0x80 (128) byte block for each URL entry until EOF
			if (stream.Read(head, 0, BlockSize) > 0)
			{
				string signature = ASCIIEncoding.ASCII.GetString(head, 0, 4);

				// validate header starts with "URL "
				if (signature.Equals(RecordSignature))
				{
					// determine record length: block count
					// this is the 4-byte int immediately after "URL ", offset is 0x04 (4)
					int numBlocks = BitConverter.ToInt32(head, RecordBlocks);
					if (numBlocks == 1)
					{
						// we already have the block
						buffer = head;
					}
					else
					{
						// we need to read numBlocks-1 since we already have the first
						buffer = new byte[BlockSize * numBlocks];
						head.CopyTo(buffer, 0);

						// read remaining blocks
						for (int i = 1; i < numBlocks; i++)
						{
							stream.Read(buffer, BlockSize * i, BlockSize);
						}
					}

					return true;
				}
			}

			buffer = new byte[0];

			return false;
		}


		//========================================================================================
		// Open()
		//========================================================================================

		/// <summary>
		/// Initialize for reading.
		/// </summary>

		public bool Open ()
		{
			try
			{
				// open index.dat
				stream = File.Open(path, FileMode.Open, FileAccess.ReadWrite, FileShare.ReadWrite);

				if ((header = ReadHeader(stream)) == null)
				{
					stream.Close();
					stream.Dispose();
					stream = null;

					throw new FileLoadException("Could not read file header");
				}

				if (!header.Signature.Equals(HeaderSignature))
				{
					stream.Close();
					stream.Dispose();
					stream = null;

					throw new FileLoadException("Invalid version");
				}

				if (!ReadHashTable(stream))
				{
					throw new FileLoadException("Could not read hash table");
				}

				// TODO: it doesn't feel right to do this here...
				rogues = new StringCollection();
				string[] names = Directory.GetFiles(root, "*.txt", SearchOption.AllDirectories);
				foreach (string name in names)
				{
					rogues.Add(Path.GetFileNameWithoutExtension(name));
				}
			}
			catch
			{
				if (stream != null)
				{
					try
					{
						stream.Close();
						stream.Dispose();
						stream = null;
					}
					catch
					{
					}
				}

				return false;
			}

			return true;
		}


		// Refer to CacheTests unit test for full explanation

		private CacheHeader ReadHeader (FileStream reader)
		{
			CacheHeader header = null;

			byte[] buffer = new byte[HeaderLen];
			if (stream.Read(buffer, 0, HeaderLen) == HeaderLen)
			{
				header = new CacheHeader(path);
				header.Signature = ASCIIEncoding.ASCII.GetString(buffer, 0, HeaderSignatureLen - 1);
				header.FileSize = BitConverter.ToInt32(buffer, HeaderFizeSizeOffset);
				header.HashTableOffset = BitConverter.ToInt32(buffer, HeaderHashOffset);
				header.TotalBlocks = BitConverter.ToInt32(buffer, HeaderBlocksOffset);
				header.AllocatedBlocks = BitConverter.ToInt32(buffer, HeaderAllocBlocksOffset);
				header.CacheQuota = BitConverter.ToInt32(buffer, HeaderQuotaOffset);
				header.CacheSize = BitConverter.ToInt32(buffer, HeaderCacheSizeOffset);
			}

			return header;
		}


		// Refer to CacheTests unit test for full explanation

		private bool ReadHashTable (FileStream reader)
		{
			if (stream.Seek(HashOffset, SeekOrigin.Begin) != HashOffset)
			{
				throw new FileLoadException("Beginning of hash table not found");
			}

			hashtable = new List<HashRef>();

			byte[] buffer;
			int blocks;
			int nextHashTable = 0;

			do
			{
				if (nextHashTable > 0)
				{
					stream.Position = nextHashTable;
					header.HashTables++;
				}

				buffer = new byte[HashHeaderLen];

				if (stream.Read(buffer, 0, HashHeaderLen) != HashHeaderLen)
				{
					throw new FileLoadException("Could not read hash table header");
				}

				string signature = ASCIIEncoding.ASCII.GetString(buffer, 0, HashSignatureLen);
				if (!signature.Equals(HashSignature))
				{
					throw new FileLoadException("Hash signature not found");
				}

				blocks = BitConverter.ToInt32(buffer, 4);
				header.HashBlocks += blocks;

				nextHashTable = BitConverter.ToInt32(buffer, 8);

				// we've read 16 bytes of 0x1000 so...
				int tableLen = (blocks * BlockSize) - HashHeaderLen;

				buffer = new byte[tableLen];
				long tableAddress = stream.Position;

				if (stream.Read(buffer, 0, tableLen) != tableLen)
				{
					throw new FileLoadException("Could not read entire hash table");
				}

				int hashkey;
				int reference;
				int index = 0;

				while (index < tableLen)
				{
					hashkey = BitConverter.ToInt32(buffer, index);

					if ((hashkey & HashValid) == 0)
					{
						reference = BitConverter.ToInt32(buffer, index + 4);
						if (((reference % BlockSize) == 0) && (reference >= HashOffset))
						{
							var hash = new HashRef()
							{
								HashAddress = tableAddress + index,
								DataAddress = reference
							};

							hashtable.Add(hash);
						}
					}
					else if ((hashkey & HashREDR) == HashREDR)
					{
						reference = BitConverter.ToInt32(buffer, index + 4);
						if (((reference % BlockSize) == 0) && (reference >= HashOffset))
						{
							//redirects.Add(number);
						}
					}

					index += 8;
				}
			} while (nextHashTable > 0);

			return hashtable.Count > 0;
		}


		#region Unit testing

		/// <summary>
		/// Get the directory path of the cookie cache.
		/// </summary>

		public string _Path
		{
			get { return path; }
		}

		#endregion Unit testing
	}
}
